setwd("/Users/ben/Documents/Papers/Submitted/SRE and spread/Outputs")

source("/Users/ben/evo-dispersal/SRE_inv/SRE_inv_functions.R")

# modified: plots population density and mean trait values through space
plotter<-function(popmatrix, spX, spY, K, label){
	png(filename=paste("DemoPlot", label, ".png", sep=""), width=10, height=15, units="cm", res=300)
	density<-table(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)) # adult density through space
	density<-density/K
  density[density==0]<-NA
	dispersal<-tapply(popmatrix[,"P"], list(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)), mean) #mean trait values through space
	hval<-tapply(popmatrix[,"H"], list(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)), mean) #mean trait values through space
	par(mfrow=c(3,1))
	X<-1:spX
	Y<-1:spY
	image(X, Y, density, zlim=c(0,2), col=rev(terrain.colors(12)), main="Population density", xlab="", ylab="y-axis distance")
	image(X, Y, dispersal, zlim=c(0,1), col=rev(terrain.colors(12)), main="D values", xlab="", ylab="y-axis distance")
	image(X, Y, hval, zlim=c(0,1), col=rev(terrain.colors(12)), main="H values", xlab="x-axis distance", ylab="y-axis distance")
	dev.off()
}

#modified: mother function
mother<-function(n=20*40*10, ngens.init=100, gens.breach=1000, spY=20, steps, K=40, lambda, d.cost, mutn=0.05, plot=F, dval=99, spread=0){
	pop<-init.inds(n, 10, spY, hvar=T, dval)
	if (dval<=1) dmut<-FALSE
	else dmut<-TRUE
	maxX<-10+spread+steps+5
	hab<-init.hab(10, spY, 0)	
	nbrs<-neighbours.init(10, spY) #reveal only x up to 10
	for (i in 1:ngens.init){
		pop<-repro.disp(pop, hab, 10, spY, K, lambda, d.cost, nbrs, mutn, adap=T, dmut=dmut)
		#if (plot==T) plotter(pop, maxX, spY, K)
	}
	init.traits<-col.sample(pop, 5) #sample initial means and variances
	nbrs<-neighbours.init(maxX, spY) #reveal all neighbours
	ngens<-ngens.init+spread+gens.breach
	hab<-init.hab(maxX, spY, steps)
	st.time<-0
	for (i in (ngens.init+1):ngens){
		pop<-repro.disp(pop, hab, maxX, spY, K, lambda, d.cost, nbrs, mutn, adap=T, dmut=dmut)
		if (plot==T) plotter(pop, maxX, spY, K, label=i)
# 		if (st.time==0 & sum((pop[,"X"])==(10+spread))>0) {
# 			plotter(pop, maxX, spY, K, i)	
# 			st.time<-i
# 		}
 		if (max(pop[,"X"])==maxX | i==ngens) {
 			return()
 		}
	}
}




mother(steps=10, lambda=2, d.cost=0.1, plot=T, spread=50, gens.breach=500, ngens.init=100)
